#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*
# ==================== 配置参数 ====================
#这仨选项必须对
DATA_YAML = "data.yaml"      #data.yaml
MODEL_PATH = "yolov8n.pt"       #这是你要用的模型
IMG_SIZE = 320                  #图像分辨率
# 显卡选择 (0: RTX3050, 1: RTX4080Super)
XianKa = 1
# 类别数量配置 (1: 单类别, 2: 二分类, 3: 三分类, 5: 五分类, 10: 十分类, 20: 多分类)
NUM_CLASSES = 3
# 训练轮数配置
EPOCHS = 20
# 自定义输出文件夹名称！#*#*#*#*#*#*#*#*#必须自定义
OUTPUT_FOLDER = "out"
# 余弦退火学习率参数，不会改别改
LR0_RTX3050 = 0.01         # RTX3050初始学习率 (推荐: SGD=0.01, AdamW=0.001)
LRF_RTX3050 = 0.001        # RTX3050最终学习率倍数 (推荐: 0.001-0.01)
LR0_RTX4080 = 0.001        # RTX4080初始学习率 (推荐: SGD=0.01, AdamW=0.001)
LRF_RTX4080 = 0.01         # RTX4080最终学习率倍数 (推荐: 0.001-0.01)
# 数据增强开关，不会改别改
ENABLE_MOSAIC = False      # 马赛克增强
ENABLE_MIXUP = False       # 混合增强
ENABLE_COPY_PASTE = False  # 复制粘贴增强
ENABLE_FLIP_LR = True      # 左右镜像
ENABLE_FLIP_UD = False     # 上下镜像
ENABLE_ROTATION = False    # 旋转增强
ENABLE_SHEAR = False       # 剪切增强
ENABLE_TRANSLATE = True    # 平移增强
ENABLE_SCALE = True        # 缩放增强
ENABLE_HSV = True          # HSV色彩空间增强
ENABLE_GRAYSCALE = False   # 灰度增强
#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*
import os
import gc
import time
import torch
import psutil
import threading
from pathlib import Path
from ultralytics import YOLO
import yaml
import onnx
class GPUTrainer:
    def __init__(self, data_yaml: str = DATA_YAML, gpu_type: int = XianKa, num_classes: int = NUM_CLASSES):
        self.data_yaml = data_yaml
        self.gpu_type = gpu_type
        self.num_classes = num_classes
        self.model = None
        self.params = self.get_rtx4080_params() if gpu_type == 1 else self.get_rtx3050_params()
        self.optimize_for_classes()
    
    def get_rtx3050_params(self):
        """RTX3050显存优化参数 - 最大化内存使用"""
        device = 0 if torch.cuda.is_available() else 'cpu'
        return {
            'model': MODEL_PATH,
            'imgsz': IMG_SIZE,
            'epochs': EPOCHS,
            'batch': 16,  # 极小batch减少显存
            'device': device,
            'optimizer': 'SGD',  # SGD比AdamW显存占用更少
            'lr0': LR0_RTX3050,
            'lrf': LRF_RTX3050,
            'cos_lr': True,
            'momentum': 0.9,
            'weight_decay': 0.0005,
            'warmup_epochs': 3.0,
            'warmup_momentum': 0.8,
            'box': 7.5,
            'cls': 0.5,
            'dfl': 1.5,
            'conf': 0.25,
            'iou': 0.7,
            'max_det': 300,  # 减少最大检测数
            'save': True,
            'save_json': False,  # 禁用JSON保存
            'half': False,  # 禁用FP16节省显存
            'amp': False,  # 禁用混合精度
            'verbose': True,
            'cache': 'ram',  # 强制使用内存缓存
            'workers': 1,  # 减少worker数
            'plots': False,  # 禁用绘图节省显存
            # 数据增强参数
            'mosaic': 1.0 if ENABLE_MOSAIC else 0.0,
            'mixup': 0.1 if ENABLE_MIXUP else 0.0,
            'copy_paste': 0.1 if ENABLE_COPY_PASTE else 0.0,
            'fliplr': 0.5 if ENABLE_FLIP_LR else 0.0,
            'flipud': 0.5 if ENABLE_FLIP_UD else 0.0,
            'degrees': 10.0 if ENABLE_ROTATION else 0.0,
            'shear': 2.0 if ENABLE_SHEAR else 0.0,
            'translate': 0.1 if ENABLE_TRANSLATE else 0.0,
            'scale': 0.5 if ENABLE_SCALE else 0.0,
            'hsv_h': 0.015 if ENABLE_HSV else 0.0,
            'hsv_s': 0.7 if ENABLE_HSV else 0.0,
            'hsv_v': 0.4 if ENABLE_HSV else 0.0,
        }

    def get_rtx4080_params(self):
        """RTX4080Super高性能参数 - 充分利用显存和性能"""
        device = 0 if torch.cuda.is_available() else 'cpu'
        return {
            'model': MODEL_PATH,
            'imgsz': IMG_SIZE,
            'epochs': EPOCHS,
            'batch': 64,  # 大batch充分利用显存
            'device': device,
            'optimizer': 'AdamW',  # AdamW优化器性能更好
            'lr0': LR0_RTX4080,
            'lrf': LRF_RTX4080,
            'cos_lr': True,
            'momentum': 0.937,
            'weight_decay': 0.0005,
            'warmup_epochs': 3.0,
            'warmup_momentum': 0.8,
            'box': 7.5,
            'cls': 0.5,
            'dfl': 1.5,
            'conf': 0.001,
            'iou': 0.7,
            'max_det': 1000,  # 增加最大检测数
            'save': True,
            'save_json': True,  # 启用JSON保存
            'half': True,  # 启用FP16提升性能
            'amp': False,  # 启用混合精度
            'verbose': True,
            'cache': 'ram',  # 使用内存缓存
            'workers': 8,  # 增加worker数
            'plots': True,  # 启用绘图
            # 数据增强参数
            'mosaic': 1.0 if ENABLE_MOSAIC else 0.0,
            'mixup': 0.1 if ENABLE_MIXUP else 0.0,
            'copy_paste': 0.1 if ENABLE_COPY_PASTE else 0.0,
            'fliplr': 0.5 if ENABLE_FLIP_LR else 0.0,
            'flipud': 0.5 if ENABLE_FLIP_UD else 0.0,
            'degrees': 10.0 if ENABLE_ROTATION else 0.0,
            'shear': 2.0 if ENABLE_SHEAR else 0.0,
            'translate': 0.1 if ENABLE_TRANSLATE else 0.0,
            'scale': 0.5 if ENABLE_SCALE else 0.0,
            'hsv_h': 0.015 if ENABLE_HSV else 0.0,
            'hsv_s': 0.7 if ENABLE_HSV else 0.0,
            'hsv_v': 0.4 if ENABLE_HSV else 0.0,
        }

    def optimize_for_classes(self):
        """根据类别数量优化参数"""
        if self.num_classes == 1:
            # 单类别优化：更高置信度阈值，减少假阳性
            self.params.update({'conf': 0.5, 'cls': 1.0, 'box': 5.0})
        elif self.num_classes == 2:
            # 二分类优化：平衡精度和召回率
            self.params.update({'conf': 0.3, 'cls': 0.8, 'box': 6.0})
        elif self.num_classes <= 5:
            # 少类别优化：标准参数
            self.params.update({'conf': 0.25, 'cls': 0.5, 'box': 7.5})
        elif self.num_classes <= 10:
            # 中等类别优化：增强分类损失权重
            self.params.update({'conf': 0.2, 'cls': 1.0, 'box': 7.0})
        else:
            # 多类别优化：更低置信度，更强分类权重
            self.params.update({'conf': 0.1, 'cls': 1.5, 'box': 6.5})

    def preheat_system(self):
        print("预热")
        if torch.cuda.is_available():
            # 通过一个随机张量进行GPU预热
            torch.cuda.empty_cache()
            dummy_data = torch.randn(1, 3, IMG_SIZE, IMG_SIZE).cuda()
            with torch.no_grad():
                _ = dummy_data * 2
            del dummy_data
            torch.cuda.empty_cache()
            print("完成")
        else:
            print("跳过GPU预热")

        # 预热CPU
        dummy_data = torch.randn(1, 3, IMG_SIZE, IMG_SIZE)
        with torch.no_grad():
            _ = dummy_data * 2  # 运行一个简单的计算来让CPU活跃
        del dummy_data
        print("CPU预热完成")

    def load_model(self):
        print(f"加载模型: {self.params['model']}")
        
        # 检查模型文件是否存在，不存在则自动下载
        model_path = self.params['model']
        if not os.path.exists(model_path):
            print(f"模型文件 {model_path} 不存在，正在下载")
        
        try:
            self.model = YOLO(self.params['model'])  # 创建YOLO模型对象，会自动下载
            print("模型加载完成")
        except Exception as e:
            print(f"模型加载失败: {e}")
            print("尝试重新下载模型...")
            # 如果加载失败，删除可能损坏的文件并重试
            if os.path.exists(model_path):
                os.remove(model_path)
                print(f"已删除损坏的模型文件: {model_path}")
            self.model = YOLO(self.params['model'])  # 重新创建，会自动下载
            print("模型重新下载并加载完成")

    def load_dataset(self):
        print(f"加载数据集: {self.data_yaml}...")
        try:
            with open(self.data_yaml, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)
            print(f"数据集加载完成 - 类别数: {self.num_classes}")
        except FileNotFoundError:
            print(f"错误: 找不到数据集文件 {self.data_yaml}")
            raise
        except yaml.YAMLError as e:
            print(f"错误: YAML文件解析失败 {e}")
            raise

    def dynamic_batch_optimization(self) -> int:
        print("动态优化batch size")
        
        if self.model is None:
            print("操？")
            return 1
        
        test_batches = [64, 32, 16, 8, 4, 2, 1] if self.gpu_type == 1 else [4, 3, 2, 1]
        optimal_batch = 1  # 默认最小值
        
        for batch_size in test_batches:
            try:
                dummy_data = torch.randn(batch_size, 3, IMG_SIZE, IMG_SIZE).cuda() if torch.cuda.is_available() else torch.randn(batch_size, 3, IMG_SIZE, IMG_SIZE)
                with torch.no_grad():
                    _ = self.model(dummy_data)
                del dummy_data
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                optimal_batch = batch_size
                print(f"最优batch size: {optimal_batch}")
                break
            except RuntimeError as e:
                if "out of memory" in str(e):
                    print(f"Batch {batch_size} 显存不足")
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                    continue
                else:
                    print(f"RuntimeError: {e}")
                    break
            except Exception as e:
                print(f"错误: {e}")
                break
        
        return optimal_batch

    def train(self):
        try:
            print("开始训练")
            print(f"数据增强配置:")
            print(f"   镜像: {'开启' if ENABLE_FLIP_LR else '关闭'}")
            print(f"   HSV增强: {'开启' if ENABLE_HSV else '关闭'}")
            print(f"   灰度增强: {'开启' if ENABLE_GRAYSCALE else '关闭'}")
            print(f"   马赛克: {'开启' if ENABLE_MOSAIC else '关闭'}")
            
            gpu_name = "rtx4080" if self.gpu_type == 1 else "rtx3050"
            train_name = f"{gpu_name}_train_{self.num_classes}class"
            self.model.train(
                data=self.data_yaml,
                project=f"runs/{OUTPUT_FOLDER}",
                name=train_name,
                exist_ok=True,
                **self.params
            )
            print("训练完成！")
        except Exception as e:
            print(f"训练错误: {str(e)}")
            raise
    
    def post_training_optimization(self):
        print("操")
        try:
            gpu_name = "rtx4080" if self.gpu_type == 1 else "rtx3050"
            train_name = f"{gpu_name}_train_{self.num_classes}class"
            results_path = f"runs/{OUTPUT_FOLDER}/{train_name}"
            
            best_model_path = Path(results_path) / 'weights' / 'best.pt'
            last_model_path = Path(results_path) / 'weights' / 'last.pt'
            
            # 导出best模型
            if best_model_path.exists():
                print(f"操: {best_model_path}")
                try:
                    model = YOLO(str(best_model_path))
                    model.export(format='onnx', data=self.data_yaml)
                    print("操！")
                except Exception as e:
                    print(f"操？: {e}")
            
            # 导出last模型
            if last_model_path.exists():
                print(f"操: {last_model_path}")
                try:
                    model = YOLO(str(last_model_path))
                    model.export(format='onnx', data=self.data_yaml)
                    print("操！")
                except Exception as e:
                    print(f"操？: {e}")
            
            print("操！")
        except Exception as e:
            print(f"操？: {e}")
    
    def run(self):
        gpu_name = "RTX4080Super" if self.gpu_type == 1 else "RTX3050"
        print(f"开始{self.num_classes}类别训练流程 ({gpu_name})...")
        print(f"初始学习率 {self.params['lr0']} 最终学习率 {self.params['lr0'] * self.params['lrf']:.6f}")
        self.preheat_system() 
        self.load_model()  
        self.load_dataset()
        optimal_batch = self.dynamic_batch_optimization()
        if optimal_batch > 0:
            self.params['batch'] = optimal_batch
        self.train()
        self.post_training_optimization()


if __name__ == "__main__":
    trainer = GPUTrainer(DATA_YAML, XianKa, NUM_CLASSES)
    trainer.run()